{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as scipy_R\n",
    "from typing import Union, List\n",
    "from shapely.geometry import LineString, Polygon\n",
    "import pickle\n",
    "from PIL import Image \n",
    "# %matplotlib notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as scipy_R\n",
    "from typing import Union, List\n",
    "from shapely.geometry import LineString, Polygon\n",
    "import pickle, torch\n",
    "from pyquaternion import Quaternion\n",
    "# with open('../tmp/result/result_aux_470.pkl','rb') as f:\n",
    "#     data = pickle.load(f)\n",
    "\n",
    "def init_axes(dpi: float=100) -> plt.Axes:\n",
    "    fig = plt.figure(dpi=dpi)\n",
    "    return fig.add_axes([0, 0, 1, 1])\n",
    "\n",
    "def matplot_2d_cloud(cloud_x, cloud_y, scale=0.1, rgb='b',\n",
    "                     alpha=1, label=None, axes: plt.Axes=None) -> None:\n",
    "    \"\"\" plot 2D cloud with matplotlib\n",
    "    :param cloud_x: list of cloud x\n",
    "    :param cloud_y: list of cloud y\n",
    "    :param scale: size of cloud points, float or list of float\n",
    "    :param rgb: color used for plot\n",
    "    :param alpha: opacity of line, 0 for transparent, 1 for opaque\n",
    "    :param label: label of cloud\n",
    "    :param axes: axes handler of matplotlib\n",
    "    \"\"\"\n",
    "    if axes is None:\n",
    "        axes = init_axes(dpi=1000)\n",
    "    if label is None:\n",
    "        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none')\n",
    "    else:\n",
    "        axes.scatter(cloud_x, cloud_y, c=rgb, s=scale, alpha=alpha, edgecolors='none', label=label)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n",
    "def matplot_polygon(points_x, points_y, scale=0.1, rgb='g',\n",
    "                    alpha=1, axes:plt.Axes=None) -> None:\n",
    "    if axes is None:\n",
    "        axes = init_axes()\n",
    "    assert len(points_x) == len(points_y) and len(points_x) > 1\n",
    "    x_arr = np.append(np.array(points_x), points_x[0])\n",
    "    y_arr = np.append(np.array(points_y), points_y[0])\n",
    "    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n",
    "def matplot_lines(points_x, points_y, scale=0.1, rgb='g',\n",
    "                    alpha=1, axes:plt.Axes=None) -> None:\n",
    "    if axes is None:\n",
    "        axes = init_axes()\n",
    "    assert len(points_x) == len(points_y) and len(points_x) > 1\n",
    "    x_arr = np.array(points_x)\n",
    "    y_arr = np.array(points_y)\n",
    "    axes.plot(x_arr, y_arr, c=rgb, lw=scale, alpha=alpha)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n",
    "\n",
    "\n",
    "DEFAULT_ALL_CLOSE_PRECISION = 1.e-4\n",
    "\n",
    "def euler_to_rotation_matrix(euler_xyz: Union[np.ndarray, list], degree=False) -> np.ndarray:\n",
    "    return scipy_R.from_euler('XYZ', euler_xyz, degrees=degree).as_matrix()\n",
    "\n",
    "def bounding_box_to_points(position: list, size: list, rotation: list) -> np.ndarray:\n",
    "    \"\"\" calculate 3D bounding box's points\n",
    "    :param position: position of box\n",
    "    :param size: size of box (x, y, z)\n",
    "    :param rotation: heading of box in euler angles\n",
    "    \"\"\"\n",
    "    assert len(position) == len(size) == len(rotation) == 3\n",
    "    pos_arr = np.array(position)\n",
    "    size_arr = np.array(size)\n",
    "    pts_prefix = np.array([[-1, -1, -1], [-1, 1, -1], [1, 1, -1], [1, -1, -1],\n",
    "                           [-1, -1, 1], [-1, 1, 1], [1, 1, 1], [1, -1, 1]]) / 2\n",
    "    return euler_to_rotation_matrix(rotation).dot((pts_prefix * size_arr).T).T + pos_arr\n",
    "def matplot_bounding_box(position: list, size: list, rotation: list, scale=0.1,\n",
    "                         rgb='k', alpha=1, axes: plt.Axes=None) -> None:\n",
    "    \"\"\" plot 3D bounding box with matplotlib\n",
    "    :param position: position of box\n",
    "    :param size: size of box (x, y, z)\n",
    "    :param rotation: heading of box in euler angles\n",
    "    :param scale: line width of box\n",
    "    :param rgb: color used for plot\n",
    "    :param alpha: opacity of line, 0 for transparent, 1 for opaque\n",
    "    :param axes: axes handler of matplotlib\n",
    "    \"\"\"\n",
    "    if axes is None:\n",
    "        axes = init_axes()\n",
    "    box_pts = bounding_box_to_points(position, size, rotation)[:5, :2]\n",
    "    axes.plot(box_pts[:, 0], box_pts[:, 1], c=rgb, lw=scale, alpha=alpha)\n",
    "    axes.set_xlabel('x axis')\n",
    "    axes.set_ylabel('y axis')\n",
    "    axes.axis('equal')\n",
    "    return axes\n",
    "\n",
    "\n",
    "def get_lidar(path):\n",
    "    return np.fromfile(path,dtype=np.float32).reshape(-1,5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gt_plot(data, gt_box=False, gt_map=False, lidar=True):\n",
    "    rgb = ['r','y','b','k']\n",
    "    for  i in range(100):\n",
    "        rgb.append('k')\n",
    "    axes = init_axes(dpi=1000)\n",
    "    if gt_map:\n",
    "        for k in range(len(data['gt_maps_3d'])):\n",
    "            gt_maps = data['gt_maps_3d'][k].cpu().numpy()\n",
    "            if len(gt_maps) <=0:\n",
    "                continue\n",
    "            for i in range(int(gt_maps[:,0].max())+1):\n",
    "                gt_map = gt_maps[gt_maps[:,0]==i,1:]\n",
    "                if len(gt_map)<=0:\n",
    "                    continue\n",
    "                class_= int(gt_map[0,-1])\n",
    "                # axes = matplot_lines(gt_map[:,0], gt_map[:,1],rgb=rgb[class_],scale=0.3,axes=axes)\n",
    "                _rgb = np.random.rand(3)\n",
    "                axes = matplot_lines(gt_map[:,0], gt_map[:,1],rgb=_rgb,scale=0.3,axes=axes)\n",
    "    # lidar\n",
    "    if lidar:\n",
    "        for k in range(len(data['points'])):\n",
    "            points = data['points'][k].cpu().numpy()\n",
    "            axes = matplot_2d_cloud(points[:,0], points[:,1],scale=0.1,rgb='g',label=4,axes=axes)\n",
    "\n",
    "    # draw bbox\n",
    "    if gt_box:\n",
    "        for k in range(len(data['gt_boxes_3d'])):\n",
    "            box = data[\"gt_boxes_3d\"][k]\n",
    "            box_center = box.center.cpu().numpy().tolist()\n",
    "            box_dims = box.dims.cpu().numpy().tolist()\n",
    "            box_yaw = box.yaw.cpu().numpy().tolist()\n",
    "            for b in range(len(box_yaw)):\n",
    "                axes = matplot_bounding_box(box_center[b],box_dims[b],[0,0,box_yaw[b]],axes=axes)\n",
    "    plt.xlim(-150, 150)\n",
    "    plt.ylim(-50, 50)\n",
    "    # plt.xlim(0, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "import argparse\n",
    "import copy\n",
    "import os\n",
    "import time\n",
    "import warnings\n",
    "from os import path as osp\n",
    "import mmcv\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "from mmcv import Config, DictAction\n",
    "from mmcv.runner import get_dist_info, init_dist\n",
    "\n",
    "from mmdet import __version__ as mmdet_version\n",
    "from mmdet3d import __version__ as mmdet3d_version\n",
    "from mmdet3d.apis import init_random_seed, train_model\n",
    "from mmdet3d.datasets import build_dataset\n",
    "from mmdet3d.models import build_model\n",
    "from mmdet3d.utils import collect_env, get_root_logger\n",
    "from mmdet.apis import set_random_seed\n",
    "from mmseg import __version__ as mmseg_version\n",
    "\n",
    "try:\n",
    "    # If mmdet version > 2.20.0, setup_multi_processes would be imported and\n",
    "    # used from mmdet instead of mmdet3d.\n",
    "    from mmdet.utils import setup_multi_processes\n",
    "except ImportError:\n",
    "    from mmdet3d.utils import setup_multi_processes\n",
    "\n",
    "\n",
    "def parse_args():\n",
    "    parser = argparse.ArgumentParser(description='Train a detector')\n",
    "    parser.add_argument('config', help='train config file path')\n",
    "    parser.add_argument('--work-dir', help='the dir to save logs and models')\n",
    "    parser.add_argument(\n",
    "        '--resume-from', help='the checkpoint file to resume from')\n",
    "    parser.add_argument(\n",
    "        '--auto-resume',\n",
    "        action='store_true',\n",
    "        help='resume from the latest checkpoint automatically')\n",
    "    parser.add_argument(\n",
    "        '--no-validate',\n",
    "        action='store_true',\n",
    "        help='whether not to evaluate the checkpoint during training')\n",
    "    group_gpus = parser.add_mutually_exclusive_group()\n",
    "    group_gpus.add_argument(\n",
    "        '--gpus',\n",
    "        type=int,\n",
    "        help='(Deprecated, please use --gpu-id) number of gpus to use '\n",
    "        '(only applicable to non-distributed training)')\n",
    "    group_gpus.add_argument(\n",
    "        '--gpu-ids',\n",
    "        type=int,\n",
    "        nargs='+',\n",
    "        help='(Deprecated, please use --gpu-id) ids of gpus to use '\n",
    "        '(only applicable to non-distributed training)')\n",
    "    group_gpus.add_argument(\n",
    "        '--gpu-id',\n",
    "        type=int,\n",
    "        default=0,\n",
    "        help='number of gpus to use '\n",
    "        '(only applicable to non-distributed training)')\n",
    "    parser.add_argument('--seed', type=int, default=0, help='random seed')\n",
    "    parser.add_argument(\n",
    "        '--diff-seed',\n",
    "        action='store_true',\n",
    "        help='Whether or not set different seeds for different ranks')\n",
    "    parser.add_argument(\n",
    "        '--deterministic',\n",
    "        action='store_true',\n",
    "        help='whether to set deterministic options for CUDNN backend.')\n",
    "    parser.add_argument(\n",
    "        '--options',\n",
    "        nargs='+',\n",
    "        action=DictAction,\n",
    "        help='override some settings in the used config, the key-value pair '\n",
    "        'in xxx=yyy format will be merged into config file (deprecate), '\n",
    "        'change to --cfg-options instead.')\n",
    "    parser.add_argument(\n",
    "        '--cfg-options',\n",
    "        nargs='+',\n",
    "        action=DictAction,\n",
    "        help='override some settings in the used config, the key-value pair '\n",
    "        'in xxx=yyy format will be merged into config file. If the value to '\n",
    "        'be overwritten is a list, it should be like key=\"[a,b]\" or key=a,b '\n",
    "        'It also allows nested list/tuple values, e.g. key=\"[(a,b),(c,d)]\" '\n",
    "        'Note that the quotation marks are necessary and that no white space '\n",
    "        'is allowed.')\n",
    "    parser.add_argument(\n",
    "        '--launcher',\n",
    "        choices=['none', 'pytorch', 'slurm', 'mpi'],\n",
    "        default='none',\n",
    "        help='job launcher')\n",
    "    parser.add_argument('--local_rank', type=int, default=0)\n",
    "    parser.add_argument(\n",
    "        '--autoscale-lr',\n",
    "        action='store_true',\n",
    "        help='automatically scale lr with the number of gpus')\n",
    "    # args = parser.parse_args(args=['configs/argov2/fusion-mt-r50_obj_map_markv_aux_l2.py','--work-dir','work_dirs/fusion_data'])\n",
    "    # args = parser.parse_args(args=['configs/hjbev/fusion4d-lidar-img-mt_obj_map_cbgs.py','--work-dir','work_dirs/fusion_data'])\n",
    "    args = parser.parse_args(args=['configs/hjbev/fusion4d-lidar-img-mt_obj_map_150x50x25_080.py','--work-dir','work_dirs/fusion_data'])\n",
    "    if 'LOCAL_RANK' not in os.environ:\n",
    "        os.environ['LOCAL_RANK'] = str(args.local_rank)\n",
    "\n",
    "    if args.options and args.cfg_options:\n",
    "        raise ValueError(\n",
    "            '--options and --cfg-options cannot be both specified, '\n",
    "            '--options is deprecated in favor of --cfg-options')\n",
    "    if args.options:\n",
    "        warnings.warn('--options is deprecated in favor of --cfg-options')\n",
    "        args.cfg_options = args.options\n",
    "\n",
    "    return args\n",
    "\n",
    "args = parse_args()\n",
    "\n",
    "cfg = Config.fromfile(args.config)\n",
    "if args.cfg_options is not None:\n",
    "    cfg.merge_from_dict(args.cfg_options)\n",
    "\n",
    "# set multi-process settings\n",
    "setup_multi_processes(cfg)\n",
    "datasets = build_dataset(cfg.data.train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fusion_frame(sequ_data, key_frame: int, ff_frame: list):\n",
    "    b = {}\n",
    "    b['points'] = sequ_data['points'].data[key_frame]\n",
    "    b['gt_maps_3d'] = sequ_data['gt_maps_3d'].data[key_frame]\n",
    "    b['gt_boxes_3d'] = sequ_data['gt_bboxes_3d'].data[key_frame]\n",
    "    \n",
    "    pose = sequ_data['img_inputs'][-1]\n",
    "    cur_pose = pose[key_frame]\n",
    "    for ff in ff_frame:\n",
    "        f_pose = pose[ff]\n",
    "        lidar2ego = cur_pose.double().inverse() @ f_pose.double()\n",
    "        lidar2ego = lidar2ego.float()\n",
    "        lidar2ego_r = lidar2ego[:3,:3]\n",
    "        lidar2ego_t = lidar2ego[:3,3]\n",
    "        # points\n",
    "        points = sequ_data['points'].data[ff][0]\n",
    "        points[:,:3] = (lidar2ego_r @ points[:,:3].unsqueeze(-1)).squeeze(-1)\n",
    "        points[:,:3] += lidar2ego_t.view(-1, 3)\n",
    "        b['points'].append(points)\n",
    "        # gt_maps\n",
    "        gt_maps = sequ_data['gt_maps_3d'].data[ff][0]\n",
    "        gt_maps[:,1:4] = (lidar2ego_r @ gt_maps[:,1:4].unsqueeze(-1)).squeeze(-1)\n",
    "        gt_maps[:,1:4] += lidar2ego_t.view(-1, 3)\n",
    "        b['gt_maps_3d'].append(gt_maps)\n",
    "    return b\n",
    "def fusion_frame_1(sequ_data, key_frame: int, fusion_frame: list):\n",
    "    b = {}\n",
    "    b['points'] = [sequ_data[key_frame]['points'][0].data]\n",
    "    b['gt_maps_3d'] = [sequ_data[key_frame]['gt_maps_3d'][-1].data]\n",
    "    b['gt_boxes_3d'] = [sequ_data[key_frame]['gt_bboxes_3d'][-1].data]\n",
    "    \n",
    "    for ff in fusion_frame:\n",
    "        cur_pose = sequ_data[key_frame]['img_inputs'][0][-1]\n",
    "        # cur_pose = pose[key_frame]\n",
    "        f_pose = sequ_data[ff]['img_inputs'][0][-1]\n",
    "        lidar2ego = cur_pose.double().inverse() @ f_pose.double()\n",
    "        lidar2ego = lidar2ego.float()\n",
    "        lidar2ego_r = lidar2ego[:3,:3]\n",
    "        lidar2ego_t = lidar2ego[:3,3]\n",
    "        # points\n",
    "        points = sequ_data[ff]['points'][0].data\n",
    "        points[:,:3] = (lidar2ego_r @ points[:,:3].unsqueeze(-1)).squeeze(-1)\n",
    "        points[:,:3] += lidar2ego_t.view(-1, 3)\n",
    "        b['points'].append(points)\n",
    "        # gt_maps\n",
    "        gt_maps = sequ_data[ff]['gt_maps_3d'][0].data\n",
    "        gt_maps[:,1:4] = (lidar2ego_r @ gt_maps[:,1:4].unsqueeze(-1)).squeeze(-1)\n",
    "        gt_maps[:,1:4] += lidar2ego_t.view(-1, 3)\n",
    "        b['gt_maps_3d'].append(gt_maps)\n",
    "    return b\n",
    "def fusion_frame_2(sequ_data, key_frame: int, fusion_frame: list):\n",
    "    b = {}\n",
    "    b['points'] = [sequ_data[key_frame]['points'].data[0][0]]\n",
    "    b['gt_maps_3d'] = [sequ_data[key_frame]['gt_maps_3d'].data[0][0]]\n",
    "    b['gt_boxes_3d'] = [sequ_data[key_frame]['gt_bboxes_3d'].data[0][0]]\n",
    "    \n",
    "    for ff in fusion_frame:\n",
    "        cur_pose = sequ_data[key_frame]['img_inputs'][-1]\n",
    "        # cur_pose = pose[key_frame]\n",
    "        f_pose = sequ_data[ff]['img_inputs'][-1]\n",
    "        lidar2ego = cur_pose.double().inverse() @ f_pose.double()\n",
    "        lidar2ego = lidar2ego.float()\n",
    "        lidar2ego_r = lidar2ego[:3,:3]\n",
    "        lidar2ego_t = lidar2ego[:3,3]\n",
    "        # points\n",
    "        points = sequ_data[ff]['points'].data[0][0]\n",
    "        points[:,:3] = (lidar2ego_r @ points[:,:3].unsqueeze(-1)).squeeze(-1)\n",
    "        points[:,:3] += lidar2ego_t.view(-1, 3)\n",
    "        b['points'].append(points)\n",
    "        # gt_maps\n",
    "        gt_maps = sequ_data[ff]['gt_maps_3d'][0].data[0][0]\n",
    "        gt_maps[:,1:4] = (lidar2ego_r @ gt_maps[:,1:4].unsqueeze(-1)).squeeze(-1)\n",
    "        gt_maps[:,1:4] += lidar2ego_t.view(-1, 3)\n",
    "        b['gt_maps_3d'].append(gt_maps)\n",
    "    return b\n",
    "def fusion_frame_3(sequ_data, key_frame: int, ff_frame: list):\n",
    "    b = {}\n",
    "    b['points'] = sequ_data['points'][0].data[key_frame]\n",
    "    b['gt_maps_3d'] = sequ_data['gt_maps_3d'][0].data[key_frame]\n",
    "    b['gt_boxes_3d'] = sequ_data['gt_bboxes_3d'][0].data[key_frame]\n",
    "    \n",
    "    pose = sequ_data['img_inputs'][-1]\n",
    "    cur_pose = pose[key_frame]\n",
    "    for ff in ff_frame:\n",
    "        f_pose = pose[ff]\n",
    "        lidar2ego = cur_pose.double().inverse() @ f_pose.double()\n",
    "        lidar2ego = lidar2ego.float()\n",
    "        lidar2ego_r = lidar2ego[:3,:3]\n",
    "        lidar2ego_t = lidar2ego[:3,3]\n",
    "        # points\n",
    "        points = sequ_data['points'].data[ff][0]\n",
    "        points[:,:3] = (lidar2ego_r @ points[:,:3].unsqueeze(-1)).squeeze(-1)\n",
    "        points[:,:3] += lidar2ego_t.view(-1, 3)\n",
    "        b['points'].append(points)\n",
    "        # gt_maps\n",
    "        gt_maps = sequ_data['gt_maps_3d'].data[ff][0]\n",
    "        gt_maps[:,1:4] = (lidar2ego_r @ gt_maps[:,1:4].unsqueeze(-1)).squeeze(-1)\n",
    "        gt_maps[:,1:4] += lidar2ego_t.view(-1, 3)\n",
    "        b['gt_maps_3d'].append(gt_maps)\n",
    "    return b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gt_plot(b,gt_box=True,gt_map=True,lidar=True)\n",
    "# Generate a function to add two numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frame =  np.random.randint(len(datasets))\n",
    "ff_frame = [i+1 for i in range(0)] # 控制显示的历史帧\n",
    "# frame = 29308\n",
    "# frame = 24400 #24410\n",
    "# frame = 36578 + 20\n",
    "a = datasets.__getitem__(frame)\n",
    "print(a['timestamp'][-1].item(),frame)\n",
    "b = fusion_frame(a, -1, ff_frame)\n",
    "gt_plot(b, gt_box=True, gt_map=True, lidar=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('mm_zhl')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "3718455ed640e0f5a83cab4f6f6b904c17414180e7ea05de7fb1835bfeac121c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
